#!/usr/bin/env python
# encoding: utf-8
"""
双GPU配置工具
用于设置和优化D2Deep模型的双GPU训练环境
"""

import torch
import torch.nn as nn
import os

def setup_dual_gpu():
    """
    设置双GPU环境
    返回配置好的设备和模型包装器
    """
    if not torch.cuda.is_available():
        print("CUDA not available, using CPU")
        return torch.device('cpu'), None, False
    
    gpu_count = torch.cuda.device_count()
    print(f"Found {gpu_count} GPU(s)")
    
    # 显示GPU信息
    for i in range(gpu_count):
        gpu_name = torch.cuda.get_device_name(i)
        gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
        print(f"GPU {i}: {gpu_name} ({gpu_memory:.1f} GB)")
    
    # 设置主设备
    device = torch.device('cuda:0')
    
    # 检查是否可以使用双GPU
    use_dual_gpu = gpu_count >= 2
    
    if use_dual_gpu:
        print("✓ Dual GPU setup available")
        print("✓ Will use DataParallel for model distribution")
        
        # 设置CUDA设备可见性（可选）
        os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
        
        # 设置内存分配策略
        torch.cuda.set_device(0)
        
        return device, [0, 1], True
    else:
        print("⚠ Only single GPU available")
        return device, None, False

def wrap_model_with_dataparallel(model, device_ids=None):
    """
    使用DataParallel包装模型
    """
    if device_ids is not None and len(device_ids) > 1:
        print(f"Wrapping model with DataParallel using GPUs: {device_ids}")
        return nn.DataParallel(model, device_ids=device_ids)
    else:
        print("Using single GPU, no DataParallel needed")
        return model

def optimize_gpu_memory():
    """
    优化GPU内存使用
    """
    if torch.cuda.is_available():
        # 清理GPU缓存
        torch.cuda.empty_cache()
        
        # 设置内存分配策略
        torch.cuda.set_per_process_memory_fraction(0.9)  # 使用90%的GPU内存
        
        print("GPU memory optimized")

def get_optimal_batch_size():
    """
    根据GPU数量返回最优批次大小
    """
    if torch.cuda.device_count() >= 2:
        return 32  # 双GPU可以使用更大的批次
    else:
        return 16  # 单GPU使用较小的批次

# 使用示例
if __name__ == "__main__":
    # 设置双GPU环境
    device, device_ids, use_dual_gpu = setup_dual_gpu()
    
    # 优化GPU内存
    optimize_gpu_memory()
    
    # 获取最优批次大小
    optimal_batch_size = get_optimal_batch_size()
    print(f"Recommended batch size: {optimal_batch_size}")
    
    print("\n" + "="*50)
    print("Dual GPU Setup Complete!")
    print("="*50)


